from queue import PriorityQueue
from utils import date_util
import pickle as pkl


class TimeData(object):
    def __init__(self, start_date, end_date, data, data_for_model):
        self.start_date = start_date
        self.end_date = end_date
        self.data_for_model = data_for_model
        self.stock_prices = {}
        self.data = data
        date = start_date
        while date != end_date:
            try:
                for i in range(0, len(data[date])):
                    key = data[date].index[i]
                    if key in ["002611.XSHE", "603110.XSHG", "002558.XSHE"]:
                        continue
                    price_list = self.stock_prices.get(key, ())
                    price_list += (data[date].close[i],)
                    self.stock_prices[key] = price_list
                date = date_util.add_one_day(date)
            except KeyError:
                date = date_util.add_one_day(date)

    def get_time_data_for_model(self):
        return self.data_for_model
        
    def get_time_data(self):
        return self.start_date, self.end_date, self.stock_prices

    def get_recent_max_increase(self, now_date, date_gap, num_stock):
        pq = PriorityQueue()
        for stock_name, price_list in self.stock_prices.items():
            now_index = date_util.date_str_diff(now_date, self.start_date)
            if now_index < date_gap:
                continue
            # 最小堆
            pq.put((price_list[now_index - date_gap] - price_list[now_index], stock_name))
            if pq.qsize() > num_stock:
                pq.get()
        stock_list = []
        while not pq.empty():
            stock_list.append(pq.get()[1])
        print("新购买了如下股票：{0}".format(stock_list))
        return stock_list
